from typing import List, Dict
from datetime import datetime

import numpy as np

from vnsumscope.trader.utility import BarGenerator
from vnsumscope.trader.object import TickData, BarData

from vnsumscope.app.portfolio_strategy import StrategyTemplate, StrategyEngine


class PairTradingStrategy(StrategyTemplate):
    """"""

    author = "用Python的交易员"

    price_add = 0.01        # 委托超价
    boll_window = 20        # 布林窗口
    boll_dev = 2            # 布林宽度
    fixed_size = 1          # 交易数量
    bond_ratio = 1          # 现券腿比例
    futures_ratio = 10      # 期货腿比例

    bond_symbol = ""        # 现券代码
    futures_symbol = ""     # 期货代码
    current_spread = 0.0    # 当前价差
    boll_mid = 0.0          # 布林中轨
    boll_down = 0.0         # 布林下轨
    boll_up = 0.0           # 布林上轨

    parameters = [
        "price_add",
        "boll_window",
        "boll_dev",
        "fixed_size",
        "bond_ratio",
        "futures_ratio",
    ]
    variables = [
        "bond_symbol",
        "futures_symbol",
        "current_spread",
        "boll_mid",
        "boll_down",
        "boll_up",
    ]

    def __init__(
        self,
        strategy_engine: StrategyEngine,
        strategy_name: str,
        vt_symbols: List[str],
        setting: dict
    ):
        """"""
        super().__init__(strategy_engine, strategy_name, vt_symbols, setting)

        # 每个合约的K线合成器字典，用于实盘中将TICK合成为1分钟K线
        self.bgs: Dict[str, BarGenerator] = {}

        # 每个合约的目标仓位字典
        self.targets: Dict[str, int] = {}

        # 最新收到的TICK时间戳，用于控制K线合成
        self.last_tick_time: datetime = None

        # 已缓存记录的K线价差数据
        self.spread_count: int = 0

        # 用numpy.array来缓存价差数据，方便向量化计算
        self.spread_data: np.array = np.zeros(100)

        # 获取现券和期货代码
        self.bond_symbol, self.futures_symbol = vt_symbols

        # 初始化K线合成器，这里不需要回调函数（传空lambda）
        for vt_symbol in self.vt_symbols:
            self.targets[vt_symbol] = 0
            self.bgs[vt_symbol] = BarGenerator(lambda bar: None)

    def on_init(self):
        """
        Callback when strategy is inited.
        """
        self.write_log("策略初始化")

        self.load_bars(1)

    def on_start(self):
        """
        Callback when strategy is started.
        """
        self.write_log("策略启动")

    def on_stop(self):
        """
        Callback when strategy is stopped.
        """
        self.write_log("策略停止")

    def on_tick(self, tick: TickData):
        """
        Callback of new tick data update.
        """
        # 当有一个TICK的时间戳，分钟发生变化时，就认为到了切片时间点
        if (
            self.last_tick_time
            and self.last_tick_time.minute != tick.datetime.minute
        ):
            bars = {}
            for vt_symbol, bg in self.bgs.items():
                bars[vt_symbol] = bg.generate()
            self.on_bars(bars)

        bg: BarGenerator = self.bgs[tick.vt_symbol]
        bg.update_tick(tick)

        self.last_tick_time = tick.datetime

    def on_bars(self, bars: Dict[str, BarData]):
        """"""
        # 全撤之前的委托
        self.cancel_all()

        # 如果有K线缺失，则忽略此次推送
        if self.bond_symbol not in bars or self.futures_symbol not in bars:
            return

        # 获取各条腿的K线数据
        bond_bar = bars[self.bond_symbol]
        futures_bar = bars[self.futures_symbol]

        # 每5根K线计算一次（可以认为在5分钟级别计算）
        if (bond_bar.datetime.minute + 1) % 5:
            return

        # 计算当前价差
        self.current_spread = bond_bar.close_price - futures_bar.close_price

        # 使用numpy数组平移的方式更新数据，而不要频繁创建销毁
        # 注意尤其不要用pandas的数据结构，实盘增量更新速度很慢
        self.spread_data[:-1] = self.spread_data[1:]
        self.spread_data[-1] = self.current_spread

        # 检查缓存数据量是否满足计算布林带的需求了
        self.spread_count += 1
        if self.spread_count <= self.boll_window:
            return

        # 切片最近的价差数据窗口，用于后续计算
        buf: np.array = self.spread_data[-self.boll_window:]

        std = buf.std()
        self.boll_mid = buf.mean()
        self.boll_up = self.boll_mid + self.boll_dev * std
        self.boll_down = self.boll_mid - self.boll_dev * std

        # 基于布林带的数值，计算交易目标仓位
        bond_pos = self.get_pos(self.bond_symbol)

        # 如果当前没有仓位
        if not bond_pos:
            if self.current_spread >= self.boll_up:
                self.targets[self.bond_symbol] = -1 * self.bond_ratio
                self.targets[self.futures_symbol] = 1 * self.futures_ratio
            elif self.current_spread <= self.boll_down:
                self.targets[self.bond_symbol] = 1 * self.bond_ratio
                self.targets[self.futures_symbol] = -1 * self.futures_ratio
        # 如果当前现券已经做多
        elif bond_pos > 0:
            if self.current_spread >= self.boll_mid:
                self.targets[self.bond_symbol] = 0
                self.targets[self.futures_symbol] = 0
        # 如果当前现券已经做空
        else:
            if self.current_spread <= self.boll_mid:
                self.targets[self.bond_symbol] = 0
                self.targets[self.futures_symbol] = 0

        # 基于目标仓位和实际仓位的差值，执行买卖交易
        for vt_symbol in self.vt_symbols:
            target_pos = self.targets[vt_symbol]
            current_pos = self.get_pos(vt_symbol)

            pos_diff = target_pos - current_pos
            volume = abs(pos_diff)
            bar = bars[vt_symbol]

            if pos_diff > 0:
                price = bar.close_price + self.price_add

                if current_pos < 0:
                    self.cover(vt_symbol, price, volume)
                else:
                    self.buy(vt_symbol, price, volume)
            elif pos_diff < 0:
                price = bar.close_price - self.price_add

                if current_pos > 0:
                    self.sell(vt_symbol, price, volume)
                else:
                    self.short(vt_symbol, price, volume)

        self.put_event()


if __name__ == '__main__':
    from datetime import datetime
    from vnsumscope.app.portfolio_strategy.backtesting import BacktestingEngine
    from vnsumscope.trader.constant import Interval
    
    # 定义回测腿代码
    bond_symbol = "190215.IB"
    futures_symbol = "T2009.CFFEX"

    # 创建回测引擎对象
    engine = BacktestingEngine()

    # 设置一系列回测参数
    engine.set_parameters(
        vt_symbols=[bond_symbol, futures_symbol],
        interval=Interval.MINUTE,
        start=datetime(2020, 1, 1),
        end=datetime(2020, 5, 30),
        rates={
            bond_symbol: 0.03/10000,
            futures_symbol: 0.03/10000
        },
        slippages={
            bond_symbol: 0,
            futures_symbol: 0
        },
        sizes={
            bond_symbol: 100000,    # 现券1000万面值
            futures_symbol: 10000   # 期货100万面值
        },
        priceticks={
            bond_symbol: 0.0001,
            futures_symbol: 0.0005
        },
        capital=10_000_000,
    )

    # 设置策略的参数，不传则用默认值
    setting = {
        "boll_window": 20,
        "boll_dev": 1,
    }

    # 添加策略到回测引擎中
    engine.add_strategy(PairTradingStrategy, setting)
    engine.load_data()

    # 执行历史数据回放，生成买卖记录
    engine.run_backtesting()

    # 基于逐日盯市计算盈亏
    df = engine.calculate_result()

    # 计算统计指标并画图
    engine.calculate_statistics()
    engine.show_chart()